package com.blue.base.oauth.server.config.authorize;

import com.blue.base.oauth.common.utils.OAuthTokenUtils;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidScopeException;
import org.springframework.security.oauth2.common.exceptions.InvalidTokenException;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.OAuth2Request;
import org.springframework.security.oauth2.provider.TokenRequest;
import org.springframework.security.oauth2.provider.token.DefaultTokenServices;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.transaction.annotation.Transactional;

import java.util.Date;
import java.util.Set;
import java.util.UUID;

/**
 * token处理类，大部分源码最初都是从父类copy，自定义其中的部分代码即可
 *
 * @author liulei
 * @version 1.0
 */
@Slf4j
@AllArgsConstructor
public class MyDefaultTokenService extends DefaultTokenServices {

    private String tokenType;

    private TokenStore tokenStore;

    private TokenEnhancer accessTokenEnhancer;


    /**
     * 允许刷新refreshToken【源码也有】
     */
    private static boolean supportRefreshToken = true;

    /**
     * 刷新refreshToken时，如果不过期则重用【源码也有】
     */
    private static boolean reuseRefreshToken = true;

    /**
     * 获取令牌accessToken入口
     */
    @Override
    @Transactional
    public OAuth2AccessToken createAccessToken(OAuth2Authentication authentication) throws AuthenticationException {
        log.info("获取token入口......");
        /**
         * 判断缓存中是否有对应client_id的登录授权缓存信息[RedisTokenStore.getAccessToken(),根据client_id+username+scope形成唯一key去查询，核心方法org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator]
         */
        OAuth2AccessToken existingAccessToken = tokenStore.getAccessToken(authentication);
        OAuth2RefreshToken refreshToken = null;
        if (existingAccessToken != null) {
            if (existingAccessToken.isExpired()) {
                if (existingAccessToken.getRefreshToken() != null) {
                    refreshToken = existingAccessToken.getRefreshToken();
                    tokenStore.removeRefreshToken(refreshToken);
                }
                tokenStore.removeAccessToken(existingAccessToken);
            } else {
                tokenStore.storeAccessToken(existingAccessToken, authentication);
                // 记录每一次token生成/刷新的时间
                OAuthTokenUtils.recordAccessTokenCreate(existingAccessToken.getValue(), existingAccessToken.getExpiration());
                return existingAccessToken;
            }
        }
        //.redis中没有缓存信息，直接创建
        if (refreshToken == null) {
            refreshToken = createRefreshToken(authentication);
        } else if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
            ExpiringOAuth2RefreshToken expiring = (ExpiringOAuth2RefreshToken) refreshToken;
            if (System.currentTimeMillis() > expiring.getExpiration().getTime()) {
                refreshToken = createRefreshToken(authentication);
            }
        }
        OAuth2AccessToken accessToken = createAccessToken(true, authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        refreshToken = accessToken.getRefreshToken();
        if (refreshToken != null) {
            tokenStore.storeRefreshToken(refreshToken, authentication);
        }
        // 记录每一次token生成/刷新的时间
        OAuthTokenUtils.recordAccessTokenCreate(accessToken.getValue(), accessToken.getExpiration());
        return accessToken;
    }

    /**
     * 刷新token入口
     */
    @Override
    @Transactional(noRollbackFor = {InvalidTokenException.class, InvalidGrantException.class})
    public OAuth2AccessToken refreshAccessToken(String refreshTokenValue, TokenRequest tokenRequest)
            throws AuthenticationException {
        log.info("刷新token入口......{}", refreshTokenValue);
        if (!supportRefreshToken) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2RefreshToken refreshToken = tokenStore.readRefreshToken(refreshTokenValue);
        if (refreshToken == null) {
            throw new InvalidGrantException("Invalid refresh token: " + refreshTokenValue);
        }

        OAuth2Authentication authentication = tokenStore.readAuthenticationForRefreshToken(refreshToken);

        String clientId = authentication.getOAuth2Request().getClientId();
        if (clientId == null || !clientId.equals(tokenRequest.getClientId())) {
            throw new InvalidGrantException("Wrong client for this refresh token: " + refreshTokenValue);
        }

        tokenStore.removeAccessTokenUsingRefreshToken(refreshToken);

        if (isExpired(refreshToken)) {
            tokenStore.removeRefreshToken(refreshToken);
            throw new InvalidTokenException("Invalid refresh token (expired): " + refreshToken);
        }

        authentication = createRefreshedAuthentication(authentication, tokenRequest);

        if (!reuseRefreshToken) {
            tokenStore.removeRefreshToken(refreshToken);
            refreshToken = createRefreshToken(authentication);
        }

        OAuth2AccessToken accessToken = createAccessToken(false, authentication, refreshToken);
        tokenStore.storeAccessToken(accessToken, authentication);
        if (!reuseRefreshToken) {
            tokenStore.storeRefreshToken(accessToken.getRefreshToken(), authentication);
        }
        // 记录每一次token生成/刷新的时间
        OAuthTokenUtils.recordAccessTokenCreate(accessToken.getValue(), accessToken.getExpiration());
        return accessToken;
    }

    private OAuth2RefreshToken createRefreshToken(OAuth2Authentication authentication) {
        if (!supportRefreshToken) {
            return null;
        }
        int validitySeconds = getRefreshTokenValiditySeconds(authentication.getOAuth2Request());
        String value = UUID.randomUUID().toString();
        if (validitySeconds > 0) {
            return new DefaultExpiringOAuth2RefreshToken(value, new Date(System.currentTimeMillis()
                    + (validitySeconds * 1000L)));
        }
        return new DefaultOAuth2RefreshToken(value);
    }

    private OAuth2AccessToken createAccessToken(boolean createNewFlag, OAuth2Authentication authentication, OAuth2RefreshToken refreshToken) {

        //根据是否直接创建accessToken的标签，进一步判断在缓存中的token是否过期，如果没有则使用新的accessToken
        String accessToken = OAuthTokenUtils.getAccessTokenByFlag(tokenType, createNewFlag);

        DefaultOAuth2AccessToken token = new DefaultOAuth2AccessToken(accessToken);
        int validitySeconds = getAccessTokenValiditySeconds(authentication.getOAuth2Request());
        if (validitySeconds > 0) {
            token.setExpiration(new Date(System.currentTimeMillis() + (validitySeconds * 1000L)));
        }
        token.setRefreshToken(refreshToken);
        token.setScope(authentication.getOAuth2Request().getScope());

        return accessTokenEnhancer != null ? accessTokenEnhancer.enhance(token, authentication) : token;
    }

    private OAuth2Authentication createRefreshedAuthentication(OAuth2Authentication authentication, TokenRequest request) {
        OAuth2Authentication narrowed = authentication;
        Set<String> scope = request.getScope();
        OAuth2Request clientAuth = authentication.getOAuth2Request().refresh(request);
        if (scope != null && !scope.isEmpty()) {
            Set<String> originalScope = clientAuth.getScope();
            if (originalScope == null || !originalScope.containsAll(scope)) {
                throw new InvalidScopeException("Unable to narrow the scope of the client authentication to " + scope
                        + ".", originalScope);
            } else {
                clientAuth = clientAuth.narrowScope(scope);
            }
        }
        narrowed = new OAuth2Authentication(clientAuth, authentication.getUserAuthentication());
        return narrowed;
    }

    @Override
    public OAuth2AccessToken readAccessToken(String accessToken) {
        OAuth2AccessToken oAuth2AccessToken = tokenStore.readAccessToken(accessToken);
        // token续签处理
        OAuth2AccessToken delayAccessToken = OAuthTokenUtils.delayAccessToken(
                oAuth2AccessToken, tokenStore,
                clientAuth -> getAccessTokenValiditySeconds(clientAuth)
        );
        return delayAccessToken == null ? oAuth2AccessToken : delayAccessToken;
    }
}
